// transpose_utils.cu
#include <cuda_runtime.h>
#include <cuda_fp16.h>

__global__ void transpose_NCHW_to_NHWC_kernel(
const float* src, __half* dst,
int N, int C, int H, int W
) {
int n = blockIdx.z;
int h = blockIdx.y * blockDim.y + threadIdx.y;
int w = blockIdx.x * blockDim.x + threadIdx.x;

if (n < N && h < H && w < W) {
for (int c = 0; c < C; c++) {
int src_idx = n*C*H*W + c*H*W + h*W + w;
int dst_idx = n*H*W*C + h*W*C + w*C + c;
dst[dst_idx] = __float2half(src[src_idx]);
}
}
}

void transpose_NCHW_to_NHWC(
float* host_src, __half* device_dst,
int N, int C, int H, int W
) {
float* device_src;
cudaMalloc(&device_src, N*C*H*W*sizeof(float));
cudaMemcpy(device_src, host_src, N*C*H*W*sizeof(float), cudaMemcpyHostToDevice);

dim3 block(16, 16);
dim3 grid((W+15)/16, (H+15)/16, N);
transpose_NCHW_to_NHWC_kernel<<<grid, block>>>(device_src, device_dst, N, C, H, W);

cudaFree(device_src);
}